import PyPDF2
from reportlab.lib.pagesizes import letter
from reportlab.lib.units import inch
from reportlab.pdfgen import canvas
from reportlab.pdfbase.ttfonts import TTFont
from reportlab.pdfbase import pdfmetrics
from datetime import datetime

from typing import List
from pikepdf import Pdf, Page, Rectangle


class Watermark:
    def __init__(self,
                 skip_pages: List[int] = [],
                 ) -> None:
        self.watermark_pdf_path = 'watermark.pdf'
        self.nrow = 1
        self.ncol = 1
        self.skip_pages = skip_pages

    def generator_score_watermark(self, teacher, score, center_x=100, center_y=100) -> None:
        pdfmetrics.registerFont(TTFont('song', 'simsun.ttc'))

        c = canvas.Canvas(self.watermark_pdf_path, pagesize=(200, 200))

        c.setLineWidth(width=4)
        c.setStrokeColorRGB(0.1, 0.6, 0.8)
        c.setFillColorRGB(0.1, 0.6, 0.8)
        c.circle(x_cen=center_x, y_cen=center_y, r=60)

        c.setLineWidth(width=1)
        c.setStrokeColorRGB(0.1, 0.6, 0.8)
        c.setFillColorRGB(0.1, 0.6, 0.8)
        c.circle(x_cen=center_x, y_cen=center_y, r=55)

        c.setFont("song", 10)
        c.drawString(center_x-20, center_y+10, '成绩：{}'.format(str(score)))
        c.drawString(center_x-30, center_y-10, '阅卷人：{}'.format(teacher))
        c.drawString(center_x-20, center_y+30, '★ ★ ★')

        timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        c.setFont("song", 9)
        c.drawString(center_x-40, center_y-30, f"{timestamp}")

        c.save()

    def add_score_watermark(self, pdf_path, text_to_find, teacher, score):
        target_pdf = Pdf.open(pdf_path, allow_overwriting_input=True)
        watermark_pdf = Pdf.open(self.watermark_pdf_path)
        watermark_page = watermark_pdf.pages[0]

        for idx, target_page in enumerate(target_pdf.pages):
            if idx in self.skip_pages:
                continue

            text_instances = self.find_text_instances(target_page, text_to_find)
            for instance in text_instances:
                x, y = instance.bbox[:2]
                target_page.add_overlay(watermark_page, Rectangle(x - 100, y - 100, x, y))
        marked_file = pdf_path.replace('.pdf', '_marked.pdf')
        target_pdf.save(marked_file)

    def find_text_instances(self, page, text):
        text_instances = []
        content = page.get_text()
        for match in content.finditer(text):
            bbox = match.bbox
            text_instances.append(Rectangle(bbox[0], bbox[1], bbox[2], bbox[3]))
        return text_instances

    def make_score_watermark(self, pdf_path, teacher, score):
        self.generator_score_watermark(teacher, score)
        self.add_score_watermark(pdf_path, "\"代码_", teacher, score)
        import os
        os.remove(self.watermark_pdf_path)


if __name__ == '__main__':
    water = Watermark()
    water.make_score_watermark('./autoExam_test.pdf', '张三', 90)
